import os
import matplotlib.animation as mpa


def save_mp4(anim: mpa.FuncAnimation,
             output_file: str,
             fps: int = 20,
             dpi: int = 150):
    """
    保存动画为MP4或GIF格式文件

    :param anim (mpa.FuncAnimation) 要保存的动画对象
    :param output_file (str) 输出文件路径
    :param fps (int) 帧率，默认为20
    :param dpi (int) 图像分辨率，默认为150
    """

    # 自动选择合适的写入器，优先使用ffmpeg，如果不可用则使用pillow
    if mpa.writers.is_available('ffmpeg'):
        writer = 'ffmpeg'
    else:
        writer = 'pillow'
        output_file = output_file.replace('.mp4', '.gif')

    # 创建输出目录并清理已存在的同名文件
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    if os.path.exists(output_file):
        os.remove(output_file)

    # 保存动画文件
    anim.save(output_file, writer=writer, fps=fps, dpi=dpi)
    print(f'[save_mp4] 动画已保存 → {output_file}  （writer={writer}）')
